import os
import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader
from sklearn.metrics import classification_report, confusion_matrix

from model import build_resnet_18_model

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
DEVICE = torch.device('cpu')
TEST_DATA_PATH = '../data/test_data.npy'
TEST_LABELS_PATH = '../data/test_labels.npy'
BEST_MODEL_SAVE_PATH = '../models/best_model.pth'
LABELS = {'非病理性': 0, '哮喘': 1, '咽炎': 2, '支气管炎': 3, '百日咳': 4}


def load_test_data():
    test_data = np.load(TEST_DATA_PATH)
    test_labels = np.load(TEST_LABELS_PATH)
    # 转换数据格式
    test_data = np.transpose(test_data, (0, 3, 1, 2))
    return test_data, test_labels


def evaluate_model(test_data):
    # 假设测试时使用 ResNet18 模型，与训练时保持一致
    input_shape = test_data.shape[1:]  # (1, H, W)
    model = build_resnet_18_model(input_shape)
    model.load_state_dict(torch.load(BEST_MODEL_SAVE_PATH, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()

    test_dataset = TensorDataset(torch.tensor(test_data, dtype=torch.float32))
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    all_preds = []
    with torch.no_grad():
        for batch in test_loader:
            inputs = batch[0].to(DEVICE)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
    return np.array(all_preds)


def print_classification_report(test_labels, predicted_classes):
    label_values = list(LABELS.values())
    label_names = list(LABELS.keys())
    print(classification_report(test_labels, predicted_classes,
                                labels=label_values, target_names=label_names, zero_division=0))


def print_confusion_matrix(test_labels, predicted_classes):
    cm = confusion_matrix(test_labels, predicted_classes)
    print("Confusion Matrix:")
    print(cm)


def evaluate():
    test_data, test_labels = load_test_data()
    predicted_classes = evaluate_model(test_data)
    print_classification_report(test_labels, predicted_classes)
    print_confusion_matrix(test_labels, predicted_classes)


if __name__ == '__main__':
    evaluate()
